// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


// Adapted from Apache Arrow.

#pragma once

#include <cstdint>
#include <limits>
#include <memory>
#include <sstream>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

// TCompactProtocol requires some #defines to work right.
#define SIGNED_RIGHT_SHIFT_IS 1
#define ARITHMETIC_RIGHT_SHIFT 1
#include <thrift/TApplicationException.h>
#include <thrift/protocol/TCompactProtocol.h>
#include <thrift/transport/TBufferTransports.h>

#include <turbo/log/logging.h>

#include <pollux/dwio/parquet/writer/arrow/exception.h>
#include <pollux/dwio/parquet/writer/arrow/file_decryptor_internal.h>
#include <pollux/dwio/parquet/writer/arrow/file_encryptor_internal.h>
#include <pollux/dwio/parquet/writer/arrow/platform.h>
#include <pollux/dwio/parquet/writer/arrow/statistics.h>

#include <pollux/dwio/parquet/writer/arrow/properties.h>
#include <pollux/dwio/parquet/writer/arrow/types.h>

namespace kumo::pollux::parquet::arrow {
    // ----------------------------------------------------------------------
    // Convert Thrift enums to Parquet enums

    // Unsafe enum converters (input is not checked for validity)

    static inline Type::type FromThriftUnsafe(
        kumo::pollux::parquet::thrift::Type::type type) {
        return static_cast<Type::type>(type);
    }

    static inline ConvertedType::type FromThriftUnsafe(
        kumo::pollux::parquet::thrift::ConvertedType::type type) {
        // item 0 is NONE
        return static_cast<ConvertedType::type>(static_cast<int>(type) + 1);
    }

    static inline Repetition::type FromThriftUnsafe(
        kumo::pollux::parquet::thrift::FieldRepetitionType::type type) {
        return static_cast<Repetition::type>(type);
    }

    static inline Encoding::type FromThriftUnsafe(
        kumo::pollux::parquet::thrift::Encoding::type type) {
        return static_cast<Encoding::type>(type);
    }

    static inline PageType::type FromThriftUnsafe(
        kumo::pollux::parquet::thrift::PageType::type type) {
        return static_cast<PageType::type>(type);
    }

    static inline Compression::type FromThriftUnsafe(
        kumo::pollux::parquet::thrift::CompressionCodec::type type) {
        switch (type) {
            case kumo::pollux::parquet::thrift::CompressionCodec::UNCOMPRESSED:
                return Compression::UNCOMPRESSED;
            case kumo::pollux::parquet::thrift::CompressionCodec::SNAPPY:
                return Compression::SNAPPY;
            case kumo::pollux::parquet::thrift::CompressionCodec::GZIP:
                return Compression::GZIP;
            case kumo::pollux::parquet::thrift::CompressionCodec::LZO:
                return Compression::LZO;
            case kumo::pollux::parquet::thrift::CompressionCodec::BROTLI:
                return Compression::BROTLI;
            case kumo::pollux::parquet::thrift::CompressionCodec::LZ4:
                return Compression::LZ4_HADOOP;
            case kumo::pollux::parquet::thrift::CompressionCodec::LZ4_RAW:
                return Compression::LZ4;
            case kumo::pollux::parquet::thrift::CompressionCodec::ZSTD:
                return Compression::ZSTD;
            default:
                DKCHECK(false) << "Cannot reach here";
                return Compression::UNCOMPRESSED;
        }
    }

    static inline BoundaryOrder::type FromThriftUnsafe(
        kumo::pollux::parquet::thrift::BoundaryOrder::type type) {
        return static_cast<BoundaryOrder::type>(type);
    }

    namespace internal {
        template<typename T>
        struct ThriftEnumTypeTraits {
        };

        template<>
        struct ThriftEnumTypeTraits<::kumo::pollux::parquet::thrift::Type::type> {
            using ParquetEnum = Type;
        };

        template<>
        struct ThriftEnumTypeTraits<
                    ::kumo::pollux::parquet::thrift::ConvertedType::type> {
            using ParquetEnum = ConvertedType;
        };

        template<>
        struct ThriftEnumTypeTraits<
                    ::kumo::pollux::parquet::thrift::FieldRepetitionType::type> {
            using ParquetEnum = Repetition;
        };

        template<>
        struct ThriftEnumTypeTraits<
                    ::kumo::pollux::parquet::thrift::Encoding::type> {
            using ParquetEnum = Encoding;
        };

        template<>
        struct ThriftEnumTypeTraits<
                    ::kumo::pollux::parquet::thrift::PageType::type> {
            using ParquetEnum = PageType;
        };

        template<>
        struct ThriftEnumTypeTraits<
                    ::kumo::pollux::parquet::thrift::BoundaryOrder::type> {
            using ParquetEnum = BoundaryOrder;
        };

        // If the parquet file is corrupted it is possible the enum value decoded
        // will not be in the range of defined values, which is undefined behaviour.
        // This facility prevents this by loading the value as the underlying type
        // and checking to make sure it is in range.

        template<
            typename EnumType,
            typename EnumTypeRaw = typename std::underlying_type<EnumType>::type>
        inline static EnumTypeRaw LoadEnumRaw(const EnumType *in) {
            EnumTypeRaw raw_value;
            // Use memcpy(), as a regular cast would be undefined behaviour on invalid
            // values
            memcpy(&raw_value, in, sizeof(EnumType));
            return raw_value;
        }

        template<typename ApiType>
        struct SafeLoader {
            using ApiTypeEnum = typename ApiType::type;
            using ApiTypeRawEnum = typename std::underlying_type<ApiTypeEnum>::type;

            template<typename ThriftType>
            inline static ApiTypeRawEnum LoadRaw(const ThriftType *in) {
                static_assert(
                    sizeof(ApiTypeEnum) == sizeof(ThriftType),
                    "parquet type should always be the same size as thrift type");
                return static_cast<ApiTypeRawEnum>(LoadEnumRaw(in));
            }

            template<typename ThriftType, bool IsUnsigned = true>
            inline static ApiTypeEnum LoadChecked(
                const typename std::enable_if<IsUnsigned, ThriftType>::type *in) {
                auto raw_value = LoadRaw(in);
                if (ARROW_PREDICT_FALSE(
                    raw_value >= static_cast<ApiTypeRawEnum>(ApiType::UNDEFINED))) {
                    return ApiType::UNDEFINED;
                }
                return FromThriftUnsafe(static_cast<ThriftType>(raw_value));
            }

            template<typename ThriftType, bool IsUnsigned = false>
            inline static ApiTypeEnum LoadChecked(
                const typename std::enable_if<!IsUnsigned, ThriftType>::type *in) {
                auto raw_value = LoadRaw(in);
                if (ARROW_PREDICT_FALSE(
                    raw_value >= static_cast<ApiTypeRawEnum>(ApiType::UNDEFINED) ||
                    raw_value < 0)) {
                    return ApiType::UNDEFINED;
                }
                return FromThriftUnsafe(static_cast<ThriftType>(raw_value));
            }

            template<typename ThriftType>
            inline static ApiTypeEnum Load(const ThriftType *in) {
                return LoadChecked<ThriftType, std::is_unsigned<ApiTypeRawEnum>::value>(in);
            }
        };
    } // namespace internal

    // Safe enum loader: will check for invalid enum value before converting

    template<
        typename ThriftType,
        typename ParquetEnum =
        typename internal::ThriftEnumTypeTraits<ThriftType>::ParquetEnum>
    inline typename ParquetEnum::type LoadEnumSafe(const ThriftType *in) {
        return internal::SafeLoader<ParquetEnum>::Load(in);
    }

    inline typename Compression::type LoadEnumSafe(
        const kumo::pollux::parquet::thrift::CompressionCodec::type *in) {
        const auto raw_value = internal::LoadEnumRaw(in);
        // Check bounds manually, as Compression::type doesn't have the same values
        // as kumo::pollux::parquet::thrift::CompressionCodec.
        const auto min_value = static_cast<decltype(raw_value)>(
            kumo::pollux::parquet::thrift::CompressionCodec::UNCOMPRESSED);
        const auto max_value = static_cast<decltype(raw_value)>(
            kumo::pollux::parquet::thrift::CompressionCodec::LZ4_RAW);
        if (raw_value < min_value || raw_value > max_value) {
            return Compression::UNCOMPRESSED;
        }
        return FromThriftUnsafe(*in);
    }

    // Safe non-enum converters

    static inline AadMetadata FromThrift(
        kumo::pollux::parquet::thrift::AesGcmV1 aesGcmV1) {
        return AadMetadata{
            aesGcmV1.aad_prefix,
            aesGcmV1.aad_file_unique,
            aesGcmV1.supply_aad_prefix
        };
    }

    static inline AadMetadata FromThrift(
        kumo::pollux::parquet::thrift::AesGcmCtrV1 aesGcmCtrV1) {
        return AadMetadata{
            aesGcmCtrV1.aad_prefix,
            aesGcmCtrV1.aad_file_unique,
            aesGcmCtrV1.supply_aad_prefix
        };
    }

    static inline EncryptionAlgorithm FromThrift(
        kumo::pollux::parquet::thrift::EncryptionAlgorithm encryption) {
        EncryptionAlgorithm encryption_algorithm;

        if (encryption.__isset.AES_GCM_V1) {
            encryption_algorithm.algorithm = ParquetCipher::AES_GCM_V1;
            encryption_algorithm.aad = FromThrift(encryption.AES_GCM_V1);
        } else if (encryption.__isset.AES_GCM_CTR_V1) {
            encryption_algorithm.algorithm = ParquetCipher::AES_GCM_CTR_V1;
            encryption_algorithm.aad = FromThrift(encryption.AES_GCM_CTR_V1);
        } else {
            throw ParquetException("Unsupported algorithm");
        }
        return encryption_algorithm;
    }

    static inline SortingColumn FromThrift(
        kumo::pollux::parquet::thrift::SortingColumn thrift_sorting_column) {
        SortingColumn sorting_column;
        sorting_column.column_idx = thrift_sorting_column.column_idx;
        sorting_column.nulls_first = thrift_sorting_column.nulls_first;
        sorting_column.descending = thrift_sorting_column.descending;
        return sorting_column;
    }

    // ----------------------------------------------------------------------
    // Convert Thrift enums from Parquet enums

    static inline kumo::pollux::parquet::thrift::Type::type ToThrift(
        Type::type type) {
        return static_cast<kumo::pollux::parquet::thrift::Type::type>(type);
    }

    static inline kumo::pollux::parquet::thrift::ConvertedType::type ToThrift(
        ConvertedType::type type) {
        // item 0 is NONE
        DKCHECK_NE(type, ConvertedType::NONE);
        // it is forbidden to emit "NA" (PARQUET-1990)
        DKCHECK_NE(type, ConvertedType::NA);
        DKCHECK_NE(type, ConvertedType::UNDEFINED);
        return static_cast<kumo::pollux::parquet::thrift::ConvertedType::type>(
            static_cast<int>(type) - 1);
    }

    static inline kumo::pollux::parquet::thrift::FieldRepetitionType::type
    ToThrift(Repetition::type type) {
        return static_cast<
            kumo::pollux::parquet::thrift::FieldRepetitionType::type>(type);
    }

    static inline kumo::pollux::parquet::thrift::Encoding::type ToThrift(
        Encoding::type type) {
        return static_cast<kumo::pollux::parquet::thrift::Encoding::type>(type);
    }

    static inline kumo::pollux::parquet::thrift::CompressionCodec::type ToThrift(
        Compression::type type) {
        switch (type) {
            case Compression::UNCOMPRESSED:
                return kumo::pollux::parquet::thrift::CompressionCodec::UNCOMPRESSED;
            case Compression::SNAPPY:
                return kumo::pollux::parquet::thrift::CompressionCodec::SNAPPY;
            case Compression::GZIP:
                return kumo::pollux::parquet::thrift::CompressionCodec::GZIP;
            case Compression::LZO:
                return kumo::pollux::parquet::thrift::CompressionCodec::LZO;
            case Compression::BROTLI:
                return kumo::pollux::parquet::thrift::CompressionCodec::BROTLI;
            case Compression::LZ4:
                return kumo::pollux::parquet::thrift::CompressionCodec::LZ4_RAW;
            case Compression::LZ4_HADOOP:
                // Deprecated "LZ4" Parquet compression has Hadoop-specific framing
                return kumo::pollux::parquet::thrift::CompressionCodec::LZ4;
            case Compression::ZSTD:
                return kumo::pollux::parquet::thrift::CompressionCodec::ZSTD;
            default:
                DKCHECK(false) << "Cannot reach here";
                return kumo::pollux::parquet::thrift::CompressionCodec::UNCOMPRESSED;
        }
    }

    static inline kumo::pollux::parquet::thrift::BoundaryOrder::type ToThrift(
        BoundaryOrder::type type) {
        switch (type) {
            case BoundaryOrder::Unordered:
            case BoundaryOrder::Ascending:
            case BoundaryOrder::Descending:
                return static_cast<kumo::pollux::parquet::thrift::BoundaryOrder::type>(
                    type);
            default:
                DKCHECK(false) << "Cannot reach here";
                return kumo::pollux::parquet::thrift::BoundaryOrder::UNORDERED;
        }
    }

    static inline kumo::pollux::parquet::thrift::SortingColumn ToThrift(
        SortingColumn sorting_column) {
        kumo::pollux::parquet::thrift::SortingColumn thrift_sorting_column;
        thrift_sorting_column.column_idx = sorting_column.column_idx;
        thrift_sorting_column.descending = sorting_column.descending;
        thrift_sorting_column.nulls_first = sorting_column.nulls_first;
        return thrift_sorting_column;
    }

    static inline kumo::pollux::parquet::thrift::Statistics ToThrift(
        const EncodedStatistics &stats) {
        kumo::pollux::parquet::thrift::Statistics statistics;
        if (stats.has_min) {
            statistics.__set_min_value(stats.min());
            // If the order is SIGNED, then the old min value must be set too.
            // This for backward compatibility
            if (stats.is_signed()) {
                statistics.__set_min(stats.min());
            }
        }
        if (stats.has_max) {
            statistics.__set_max_value(stats.max());
            // If the order is SIGNED, then the old max value must be set too.
            // This for backward compatibility
            if (stats.is_signed()) {
                statistics.__set_max(stats.max());
            }
        }
        if (stats.has_null_count) {
            statistics.__set_null_count(stats.null_count);
        }
        if (stats.has_distinct_count) {
            statistics.__set_distinct_count(stats.distinct_count);
        }

        return statistics;
    }

    static inline kumo::pollux::parquet::thrift::AesGcmV1 ToAesGcmV1Thrift(
        AadMetadata aad) {
        kumo::pollux::parquet::thrift::AesGcmV1 aesGcmV1;
        // aad_file_unique is always set
        aesGcmV1.__set_aad_file_unique(aad.aad_file_unique);
        aesGcmV1.__set_supply_aad_prefix(aad.supply_aad_prefix);
        if (!aad.aad_prefix.empty()) {
            aesGcmV1.__set_aad_prefix(aad.aad_prefix);
        }
        return aesGcmV1;
    }

    static inline kumo::pollux::parquet::thrift::AesGcmCtrV1 ToAesGcmCtrV1Thrift(
        AadMetadata aad) {
        kumo::pollux::parquet::thrift::AesGcmCtrV1 aesGcmCtrV1;
        // aad_file_unique is always set
        aesGcmCtrV1.__set_aad_file_unique(aad.aad_file_unique);
        aesGcmCtrV1.__set_supply_aad_prefix(aad.supply_aad_prefix);
        if (!aad.aad_prefix.empty()) {
            aesGcmCtrV1.__set_aad_prefix(aad.aad_prefix);
        }
        return aesGcmCtrV1;
    }

    static inline kumo::pollux::parquet::thrift::EncryptionAlgorithm ToThrift(
        EncryptionAlgorithm encryption) {
        kumo::pollux::parquet::thrift::EncryptionAlgorithm encryption_algorithm;
        if (encryption.algorithm == ParquetCipher::AES_GCM_V1) {
            encryption_algorithm.__set_AES_GCM_V1(ToAesGcmV1Thrift(encryption.aad));
        } else {
            encryption_algorithm.__set_AES_GCM_CTR_V1(
                ToAesGcmCtrV1Thrift(encryption.aad));
        }
        return encryption_algorithm;
    }

    // ----------------------------------------------------------------------
    // Thrift struct serialization / deserialization utilities

    using ThriftBuffer = apache::thrift::transport::TMemoryBuffer;

    class ThriftDeserializer {
    public:
        explicit ThriftDeserializer(const ReaderProperties &properties)
            : ThriftDeserializer(
                properties.thrift_string_size_limit(),
                properties.thrift_container_size_limit()) {
        }

        ThriftDeserializer(int32_t string_size_limit, int32_t container_size_limit)
            : string_size_limit_(string_size_limit),
              container_size_limit_(container_size_limit) {
        }

        // Deserialize a thrift message from buf/len.  buf/len must at least contain
        // all the bytes needed to store the thrift message.  On return, len will be
        // set to the actual length of the header.
        template<class T>
        void DeserializeMessage(
            const uint8_t *buf,
            uint32_t *len,
            T *deserialized_msg,
            const std::shared_ptr<Decryptor> &decryptor = NULLPTR) {
            if (decryptor == NULLPTR) {
                // thrift message is not encrypted
                DeserializeUnencryptedMessage(buf, len, deserialized_msg);
            } else {
                // thrift message is encrypted
                uint32_t clen;
                clen = *len;
                // decrypt
                auto decrypted_buffer =
                        std::static_pointer_cast<ResizableBuffer>(AllocateBuffer(
                            decryptor->pool(),
                            static_cast<int64_t>(clen - decryptor->CiphertextSizeDelta())));
                const uint8_t *cipher_buf = buf;
                uint32_t decrypted_buffer_len =
                        decryptor->Decrypt(cipher_buf, 0, decrypted_buffer->mutable_data());
                if (decrypted_buffer_len <= 0) {
                    throw ParquetException("Couldn't decrypt buffer\n");
                }
                *len = decrypted_buffer_len + decryptor->CiphertextSizeDelta();
                DeserializeUnencryptedMessage(
                    decrypted_buffer->data(), &decrypted_buffer_len, deserialized_msg);
            }
        }

    private:
        // On Thrift 0.14.0+, we want to use TConfiguration to raise the max message
        // size limit (ARROW-13655).  If we wanted to protect against huge messages,
        // we could do it ourselves since we know the message size up front.
        std::shared_ptr<ThriftBuffer> CreateReadOnlyMemoryBuffer(
            uint8_t *buf,
            uint32_t len) {
#if PARQUET_THRIFT_VERSION_MAJOR > 0 || PARQUET_THRIFT_VERSION_MINOR >= 14
    auto conf = std::make_shared<apache::thrift::TConfiguration>();
    conf->setMaxMessageSize(std::numeric_limits<int>::max());
    return std::make_shared<ThriftBuffer>(
        buf, len, ThriftBuffer::OBSERVE, conf);
#else
            return std::make_shared<ThriftBuffer>(buf, len);
#endif
        }

        template<class T>
        void DeserializeUnencryptedMessage(
            const uint8_t *buf,
            uint32_t *len,
            T *deserialized_msg) {
            // Deserialize msg bytes into c++ thrift msg using memory transport.
            auto tmem_transport =
                    CreateReadOnlyMemoryBuffer(const_cast<uint8_t *>(buf), *len);
            apache::thrift::protocol::TCompactProtocolFactoryT<ThriftBuffer>
                    tproto_factory;
            // Protect against CPU and memory bombs
            tproto_factory.setStringSizeLimit(string_size_limit_);
            tproto_factory.setContainerSizeLimit(container_size_limit_);
            auto tproto = tproto_factory.getProtocol(tmem_transport);
            try {
                deserialized_msg->read(tproto.get());
            } catch (std::exception &e) {
                std::stringstream ss;
                ss << "Couldn't deserialize thrift: " << e.what() << "\n";
                throw ParquetException(ss.str());
            }
            uint32_t bytes_left = tmem_transport->available_read();
            *len = *len - bytes_left;
        }

        const int32_t string_size_limit_;
        const int32_t container_size_limit_;
    };

    /// Utility class to serialize thrift objects to a binary format.  This object
    /// should be reused if possible to reuse the underlying memory.
    /// Note: thrift will encode NULLs into the serialized buffer so it is not valid
    /// to treat it as a string.
    class ThriftSerializer {
    public:
        explicit ThriftSerializer(int initial_buffer_size = 1024)
            : mem_buffer_(new ThriftBuffer(initial_buffer_size)) {
            apache::thrift::protocol::TCompactProtocolFactoryT<ThriftBuffer> factory;
            protocol_ = factory.getProtocol(mem_buffer_);
        }

        /// Serialize obj into a memory buffer.  The result is returned in buffer/len.
        /// The memory returned is owned by this object and will be invalid when
        /// another object is serialized.
        template<class T>
        void SerializeToBuffer(const T *obj, uint32_t *len, uint8_t **buffer) {
            SerializeObject(obj);
            mem_buffer_->getBuffer(buffer, len);
        }

        template<class T>
        void SerializeToString(const T *obj, std::string *result) {
            SerializeObject(obj);
            *result = mem_buffer_->getBufferAsString();
        }

        template<class T>
        int64_t Serialize(
            const T *obj,
            ArrowOutputStream *out,
            const std::shared_ptr<Encryptor> &encryptor = NULLPTR) {
            uint8_t *out_buffer;
            uint32_t out_length;
            SerializeToBuffer(obj, &out_length, &out_buffer);

            // obj is not encrypted
            if (encryptor == NULLPTR) {
                PARQUET_THROW_NOT_OK(out->Write(out_buffer, out_length));
                return static_cast<int64_t>(out_length);
            } else {
                // obj is encrypted
                return SerializeEncryptedObj(out, out_buffer, out_length, encryptor);
            }
        }

    private:
        template<class T>
        void SerializeObject(const T *obj) {
            try {
                mem_buffer_->resetBuffer();
                obj->write(protocol_.get());
            } catch (std::exception &e) {
                std::stringstream ss;
                ss << "Couldn't serialize thrift: " << e.what() << "\n";
                throw ParquetException(ss.str());
            }
        }

        int64_t SerializeEncryptedObj(
            ArrowOutputStream *out,
            uint8_t *out_buffer,
            uint32_t out_length,
            const std::shared_ptr<Encryptor> &encryptor) {
            auto cipher_buffer =
                    std::static_pointer_cast<ResizableBuffer>(AllocateBuffer(
                        encryptor->pool(),
                        static_cast<int64_t>(
                            encryptor->CiphertextSizeDelta() + out_length)));
            int cipher_buffer_len = encryptor->Encrypt(
                out_buffer, out_length, cipher_buffer->mutable_data());

            PARQUET_THROW_NOT_OK(out->Write(cipher_buffer->data(), cipher_buffer_len));
            return static_cast<int64_t>(cipher_buffer_len);
        }

        std::shared_ptr<ThriftBuffer> mem_buffer_;
        std::shared_ptr<apache::thrift::protocol::TProtocol> protocol_;
    };
} // namespace kumo::pollux::parquet::arrow
